/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define SPYDLL_BUILD_DLL

#define TRACE_LEVEL	2 /* warning */


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "HookSys.h"

// standard headers
#include <tchar.h>
#include <shlwapi.h> // functions PathXxx

// project's headers
#include "Assert.h"
#include "SpyDll.h"
#include "SpyDllInt.h"
#include "Hooks.h"
#include "Strlcpy.h"
#include "Trace.h"


/******************************************************************/
/* Internal functions                                             */
/******************************************************************/

BOOL _HookSys_ExtractPathFromCmdLine (LPTSTR szPath, LPCTSTR szCurDir, LPCTSTR szCmdLine, 
				      BOOL bRemoveArgs, BOOL bSearchPath) ;


/******************************************************************/
/* Exported function : CreateProcessA                             */
/******************************************************************/

BOOL WINAPI Spy_CreateProcessA (LPCSTR			szAppName,
				LPSTR			szCmdLine,
				LPSECURITY_ATTRIBUTES	pProcessAttr,
				LPSECURITY_ATTRIBUTES	pThreadAttr,
				BOOL			bInheritHandles,
				DWORD			dwCreationFlags,
				PVOID			pEnvironment,
				LPCSTR			szCurrentDirectory,
				LPSTARTUPINFOW		pStartupInfo,
				LPPROCESS_INFORMATION	pProcessInformation) 
{
  PROC			pfnStub = Hooks_GetStubAddress (HOOKS_CREATEPROCESSA) ;

  WORD			nReaction ;
  BOOL			bSuccess ;
  DWORD			dwServerResponse ;
  DWORD			dwSaveErrorCode ;
  SDNCREATINGPROCESS	sdncp ;
  SDNPROCESSCREATED	sdnpc ;
  PROCESS_INFORMATION	pi ;

  BOOL			bAppNameSpecified ;
  LPTSTR		tszCmdLine, tszCurDir ;
  TCHAR			tszPath[MAX_PATH] ;

  TRACE ;

  if( szAppName==NULL && szCmdLine==NULL )
    TRACE_WARNING (TEXT("szAppName==NULL and szCmdLine==NULL\n"));

  bAppNameSpecified = szAppName != NULL ;
  
  // get command line in unicode format
  tszCmdLine = SpyInt_AsciiToUnicode (bAppNameSpecified ? szAppName : szCmdLine) ;

  // get directory in unicode format
  tszCurDir = SpyInt_AsciiToUnicode (szCurrentDirectory) ;
   
  // extract program path 
  _HookSys_ExtractPathFromCmdLine (tszPath, tszCurDir, tszCmdLine, 
				   !bAppNameSpecified, !bAppNameSpecified) ;

  TRACE_INFO (TEXT("bAppNameSpecified = %d\n"), bAppNameSpecified) ;
  TRACE_INFO (TEXT("CmdLine = %s\n"), tszCmdLine) ;
  TRACE_INFO (TEXT("Path = %s\n"), tszPath) ; 

  // apply filter
  nReaction = SpyInt_CatchCall (REASON_SYS_EXECUTE, TEXT("s"), tszPath) ;

  // if filter reject call, just return false
  if( nReaction == RULE_REJECT ) {
    free (tszCurDir) ;
    free (tszCmdLine) ;
    SetLastError (ERROR_ACCESS_DENIED) ;
    return FALSE ; 
  }

  // if filter reject call, just return true
  if( nReaction == RULE_FEIGN ) 
    {
      free (tszCurDir) ;
      free (tszCmdLine) ;
      
      if( pProcessInformation ) {
	pProcessInformation->hProcess = NULL ;
	pProcessInformation->hThread = NULL ;
      }

      SetLastError (ERROR_SUCCESS) ;
      return TRUE ; 
    }

  // Notify server that a process will be created
  // so it will not try to inject it in another manner.
  // The server response is 1 if we are allowed to hook.
  sdncp.hdr.dwCode	= SDN_CREATING_PROCESS ;
  sdncp.hdr.dwProcessId = GetCurrentProcessId () ;
  _tcslcpy (sdncp.szPath, tszPath, MAX_PATH) ;
  dwServerResponse = SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
  TRACE_INFO (TEXT("Server response : %u\n"), dwServerResponse) ;
  
  // provide a local PROCESS_INFORMATION if needed
  if( ! pProcessInformation )
    pProcessInformation = &pi ;

  // call real CreateProcessA
  bSuccess = pfnStub (szAppName, szCmdLine, pProcessAttr, pThreadAttr,
		      bInheritHandles, dwCreationFlags|CREATE_SUSPENDED, 
		      pEnvironment, szCurrentDirectory, pStartupInfo, 
		      pProcessInformation) ;
  
  // notify server if failed
  if( ! bSuccess )
    {
      dwSaveErrorCode = GetLastError () ;
      sdncp.hdr.dwCode = SDN_CREATING_FAILED ;
      SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
      free (tszCurDir) ;
      free (tszCmdLine) ;
      SetLastError (dwSaveErrorCode) ;
      return FALSE ;
    }

  // inject process if allowed
  if( dwServerResponse==1 )
    {
      if( ! SpyDll_InjectInSuspendedThread(pProcessInformation->hProcess, 
					   pProcessInformation->hThread) )
	TRACE_ERROR (TEXT("InjectInSuspendedThread failed\n")) ;     	  
    }

  // resume process if needed
  if( ! (dwCreationFlags&CREATE_SUSPENDED) ) 
    if( 1!=ResumeThread (pProcessInformation->hThread) )
      TRACE_ERROR (TEXT("Failed to resume thread")) ;
  
  // notify server
  sdnpc.hdr.dwCode	= SDN_PROCESS_CREATED ;
  sdnpc.hdr.dwProcessId	= GetCurrentProcessId () ;
  sdnpc.bHooked		= 0!=dwServerResponse ;
  sdnpc.dwProcessId	= pProcessInformation->dwProcessId ;
  _tcslcpy (sdnpc.szPath, tszPath, MAX_PATH) ;
  SpyInt_NotifyServer (&sdnpc.hdr, sizeof(sdnpc)) ;

  // now the unicode strings can be freed		  
  free (tszCurDir) ;
  free (tszCmdLine) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function : CreateProcessW                             */
/******************************************************************/

BOOL WINAPI Spy_CreateProcessW (LPCWSTR			szAppName,
				LPWSTR			szCmdLine,
				LPSECURITY_ATTRIBUTES	pProcessAttr,
				LPSECURITY_ATTRIBUTES	pThreadAttr,
				BOOL			bInheritHandles,
				DWORD			dwCreationFlags,
				PVOID			pEnvironment,
				LPCWSTR			szCurrentDirectory,
				LPSTARTUPINFOW		pStartupInfo,
				LPPROCESS_INFORMATION	pProcessInformation) 
{
  PROC			pfnStub = Hooks_GetStubAddress (HOOKS_CREATEPROCESSW) ;

  DWORD			nReaction ;
  BOOL			bSuccess ;
  DWORD			dwServerResponse ;
  DWORD			dwSaveErrorCode ;
  SDNCREATINGPROCESS	sdncp ;
  SDNPROCESSCREATED	sdnpc ;
  PROCESS_INFORMATION	pi ;

  BOOL			bAppNameSpecified ;
  LPCTSTR		tszCmdLine ;
  TCHAR			tszPath[MAX_PATH] ;

  TRACE ;

  if( szAppName==NULL && szCmdLine==NULL )
    TRACE_WARNING (TEXT("szAppName==NULL and szCmdLine==NULL\n"));

  bAppNameSpecified = szAppName != NULL ;
  
  // get command line
  tszCmdLine = bAppNameSpecified ? szAppName : szCmdLine ;

  // extract program path 
  _HookSys_ExtractPathFromCmdLine (tszPath, szCurrentDirectory, tszCmdLine, 
				   !bAppNameSpecified, !bAppNameSpecified) ;

  TRACE_INFO (TEXT("bAppNameSpecified = %d\n"), bAppNameSpecified) ;
  TRACE_INFO (TEXT("CmdLine = %s\n"), tszCmdLine) ;
  TRACE_INFO (TEXT("Path = %s\n"), tszPath) ; 

  // if the file doesn't exists, don't notify the server
  if( 0xFFFFFFFF==GetFileAttributes(tszPath) )
    {
      SetLastError (ERROR_FILE_NOT_FOUND) ;
      return FALSE ; 
    }

  // apply filter
  nReaction = SpyInt_CatchCall (REASON_SYS_EXECUTE, TEXT("s"), tszPath) ;

  // if filter reject call, just return false
  if( nReaction == RULE_REJECT ) {
    SetLastError (ERROR_ACCESS_DENIED) ;
    return FALSE ; 
  }

  // if filter reject call, just return true
  if( nReaction == RULE_FEIGN ) 
    {
      if( pProcessInformation ) {
	pProcessInformation->hProcess = NULL ;
	pProcessInformation->hThread = NULL ;
      }

      SetLastError (ERROR_SUCCESS) ;
      return TRUE ;       
    }

  // Notify server that a process will be created
  // so it will not try to inject it in another manner.
  // The server response is 1 if we are allowed to hook.
  sdncp.hdr.dwCode	= SDN_CREATING_PROCESS ;
  sdncp.hdr.dwProcessId = GetCurrentProcessId () ;
  _tcslcpy (sdncp.szPath, tszPath, MAX_PATH) ;
  dwServerResponse = SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
  TRACE_INFO (TEXT("Server response : %u\n"), dwServerResponse) ;
 
  // provide a local PROCESS_INFORMATION if needed
  if( ! pProcessInformation )
    pProcessInformation = &pi ;

  // call real CreateProcessW
  bSuccess = pfnStub (szAppName, szCmdLine, pProcessAttr, pThreadAttr,
		      bInheritHandles, dwCreationFlags|CREATE_SUSPENDED, 
		      pEnvironment, szCurrentDirectory, pStartupInfo, 
		      pProcessInformation) ;
  
  // notify server if failed
  if( ! bSuccess )
    {
      dwSaveErrorCode = GetLastError () ;
      sdncp.hdr.dwCode = SDN_CREATING_FAILED ;
      SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
      SetLastError (dwSaveErrorCode) ;
      return FALSE ;
    }

  // inject process if allowed
  if( dwServerResponse==1 )
    {
      if( ! SpyDll_InjectInSuspendedThread(pProcessInformation->hProcess, 
					   pProcessInformation->hThread) )
	TRACE_ERROR (TEXT("InjectInSuspendedThread failed\n")) ;
    }
  
  // resume process if needed
  if( ! (dwCreationFlags&CREATE_SUSPENDED) ) 
    if( 1!=ResumeThread (pProcessInformation->hThread) )
      TRACE_ERROR (TEXT("Failed to resume thread")) ;
  
  // notify server
  sdnpc.hdr.dwCode	= SDN_PROCESS_CREATED ;
  sdnpc.hdr.dwProcessId	= GetCurrentProcessId () ;
  sdnpc.dwProcessId	= pProcessInformation->dwProcessId ;
  sdnpc.bHooked		= 0!=dwServerResponse ;
  _tcslcpy (sdnpc.szPath, tszPath, MAX_PATH) ;
  SpyInt_NotifyServer (&sdnpc.hdr, sizeof(sdnpc)) ;
   
  return bSuccess ;
}


/******************************************************************/
/* Internal function : ExtractPathFromCmdLine                     */
/******************************************************************/

BOOL _HookSys_ExtractPathFromCmdLine (LPTSTR szPath, LPCTSTR szCurDir, LPCTSTR szCmdLine, 
				      BOOL bRemoveArgs, BOOL bSearchPath)
{
  TCHAR szAppName[MAX_PATH] ;

  if( bRemoveArgs )
    {
      int iWrite, iRead ;
      BOOL bStopOnQuote ;
      
      bStopOnQuote = szCmdLine[0]==TEXT('\"') ;
      
      iWrite = 0 ;
      for( iRead=bStopOnQuote ; szCmdLine[iRead] ; iRead++ )
	{
	  if( szCmdLine[iRead]==TEXT('\"') && bStopOnQuote )
	    break ;

	  if( szCmdLine[iRead]==TEXT(' ') && ! bStopOnQuote )
	    {
	      szAppName[iWrite] = 0 ;
	      
	      if( 0<SearchPath(NULL,szAppName,TEXT(".exe"),MAX_PATH,szPath,NULL) )
		return TRUE ;

	      TRACE_WARNING (TEXT("File \"%s\" doesn't exist\n"), szAppName) ;
	    }

	  szAppName[iWrite++] = szCmdLine[iRead] ;

	  if( iWrite>=MAX_PATH ) break ;
	}

      szAppName[iWrite] = 0 ;
   
      if( ! SearchPath(NULL,szAppName,TEXT(".exe"),MAX_PATH,szPath,NULL) )
	_tcscpy (szPath, szAppName) ;
    }
  else
    {
      if( ! SearchPath(NULL,szAppName,TEXT(".exe"),MAX_PATH,szPath,NULL) )
	_tcscpy (szPath, szCmdLine) ;   
    }

  return TRUE ;
}

